from .nodes.model_sampling import ModelSamplingFluxNormalized
from .nodes.latent_size_picker import FluxEmptyLatentSizePicker
from .nodes.soul_samplers import SoulSampler, SoulSamplerAdvanced, SoulSamplerHybrid, SoulSamplerHybridAdvanced, SoulSamplerDPM, SoulSamplerDPMAdvanced
from .nodes.highresfix_injection import HighResFixApply, HighResFixModelInjection, HighResFixConditioningDuplicator
from .nodes.stork import STORKSamplerSelect

NODE_CLASS_MAPPINGS = {
    "ModelSamplingFluxNormalized": ModelSamplingFluxNormalized,
    "FluxEmptyLatentSizePicker": FluxEmptyLatentSizePicker,
    "SoulSampler": SoulSampler,
    "SoulSamplerAdvanced": SoulSamplerAdvanced,
    "SoulSamplerHybrid": SoulSamplerHybrid,
    "SoulSamplerHybridAdvanced": SoulSamplerHybridAdvanced,
    "SoulSamplerDPM": SoulSamplerDPM,
    "SoulSamplerDPMAdvanced": SoulSamplerDPMAdvanced,
    "HighResFixApply": HighResFixApply,
    "HighResFixModelInjection": HighResFixModelInjection,
    "HighResFixConditioningDuplicator": HighResFixConditioningDuplicator,
    "STORKSamplerSelect": STORKSamplerSelect,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ModelSamplingFluxNormalized": "Model Sampling Flux Normalized",
    "FluxEmptyLatentSizePicker": "Flux Empty Latent Size Picker",
    "SoulSampler": "Soul Sampler",
    "SoulSamplerAdvanced": "Soul Sampler (Advanced)",
    "SoulSamplerHybrid": "Soul Sampler Hybrid",
    "SoulSamplerHybridAdvanced": "Soul Sampler Hybrid (Advanced)",
    "SoulSamplerDPM": "Soul Sampler DPM++",
    "SoulSamplerDPMAdvanced": "Soul Sampler DPM++ (Advanced)",
    "HighResFixApply": "HighResFix Apply",
    "HighResFixModelInjection": "HighResFix Model Injection (Advanced)",
    "HighResFixConditioningDuplicator": "HighResFix Conditioning Duplicator (Advanced)",
    "STORKSamplerSelect": "STORK Sampler (42lux)",
}

__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']